import requests
from bs4 import BeautifulSoup
import random


def get_media_src(parent_tag):
    """
    从父级div中提取视频或图片链接。
    优先取 <video src>，否则取 <img src> 或 <img data-src>。
    """
    # 先找 video
    video_tag = parent_tag.find("video")
    if video_tag and video_tag.get("src"):
        return {
            "src": video_tag["src"],
            "type": "video",
        }

    # 再找 img
    img_tag = parent_tag.find("img")
    if img_tag:
        return {
            "src": img_tag["src"],
            "type": "image",
        }

    # 如果都没有，返回 None
    return None


def parse_items(page):
    if page == 1:
        url = "https://haowallpaper.com/homeView"
    else:
        url = "https://haowallpaper.com/homeView?isSel=true&page=" + str(page)
    response = requests.get(url)
    # print(response.text)
    soup = BeautifulSoup(response.text, "html.parser")
    soup_find_all = soup.find_all('div', class_='card')
    items = []
    for item in soup_find_all:
        items.append(get_media_src(item))
    return items


def download_file(item):
    """
    根据链接下载文件并保存
    :param url: 文件的下载链接
    :param filename: 本地保存的文件名
    """
    response = requests.get(item["src"], stream=True)
    response.raise_for_status()  # 如果请求失败，抛出异常

    filename = fr"C:\Users\Administrator\Desktop\wallpaper\{item["src"].split('/')[-1]}.{"mp4" if item['type'] == "video" else "png"}"

    with open(filename, "wb") as f:
        for chunk in response.iter_content(chunk_size=8192):
            if chunk:
                f.write(chunk)

    print(f"下载完成: {filename}")


if __name__ == '__main__':
    num = random.randint(1, 2019)
    items = parse_items(num)
    print(items)
    print(f"共有{len(items)}条数据")
    while True:
        try:
            download_num = int(input("请输入想要下载的下标数："))
            if 0 < download_num <= len(items):
                break
            else:
                print(f"❌ 输入必须是 1 ~ {len(items)} 的正整数")
        except ValueError:
            print("❌ 输入无效，请输入一个正整数")

    # 下载对应文件
    download_file(items[download_num - 1])
